// TensorFlow.js model initialization and training logic
let tfModel = null;
let tfTrainingHistory = {
    accuracy: [],
    valAccuracy: [],
    loss: [],
    valLoss: []
};

/**
 * Initialize a fraud detection model using TensorFlow.js
 * Architecture: Sequential model with dense layers for binary classification
 */
function initializeModel(inputFeatures = 30) {
    tfModel = tf.sequential({
        layers: [
            // Input layer + first hidden layer with ReLU activation
            tf.layers.dense({
                inputShape: [inputFeatures],
                units: 64,
                activation: 'relu',
                kernelInitializer: 'glorotNormal'
            }),
            // Dropout layer for regularization
            tf.layers.dropout({ rate: 0.3 }),
            // Second hidden layer
            tf.layers.dense({
                units: 32,
                activation: 'relu',
                kernelInitializer: 'glorotNormal'
            }),
            // Output layer with sigmoid activation for binary classification
            tf.layers.dense({
                units: 1,
                activation: 'sigmoid'
            })
        ]
    });

    // Compile the model with Adam optimizer and binary crossentropy loss
    tfModel.compile({
        optimizer: tf.train.adam(parseFloat(document.getElementById('learning-rate').value) || 0.005),
        loss: 'binaryCrossentropy',
        metrics: ['accuracy']
    });

    console.log('TensorFlow.js model initialized successfully');
    return tfModel;
}

/**
 * Prepare data for training using TensorFlow.js tensors
 * @param {Array} data - Raw transaction data
 * @returns {Object} - Tensor objects for features and labels
 */
function prepareTensors(data) {
    // Extract features (V1-V28, Time, Amount) and labels (Class)
    const features = data.map(d => [
        d.Time, d.V1, d.V2, d.V3, d.V4, d.V5, d.V6, d.V7, d.V8, d.V9, d.V10,
        d.V11, d.V12, d.V13, d.V14, d.V15, d.V16, d.V17, d.V18, d.V19, d.V20,
        d.V21, d.V22, d.V23, d.V24, d.V25, d.V26, d.V27, d.V28, d.Amount
    ]);
    const labels = data.map(d => [d.Class]);

    // Convert to tensors and normalize features
    const xs = tf.tensor2d(features).div(tf.scalar(1000)); // Simple normalization
    const ys = tf.tensor2d(labels);

    return { xs, ys };
}

/**
 * Train the TensorFlow.js model with processed data
 * @param {Array} trainingData - Labeled transaction data
 * @param {number} epochs - Number of training epochs
 * @param {number} batchSize - Batch size for training
 */
async function trainTensorFlowModel(trainingData, epochs = 5, batchSize = 2048) {
    if (!trainingData || trainingData.length === 0) {
        alert('No training data available. Upload and process data first.');
        return;
    }

    // Initialize model if not already created
    if (!tfModel) initializeModel();

    // Prepare tensors from training data
    const { xs, ys } = prepareTensors(trainingData);

    // Split data into training (80%) and validation (20%) sets
    const splitIndex = Math.floor(trainingData.length * 0.8);
    const [trainXs, valXs] = [
        xs.slice([0, 0], [splitIndex, -1]),
        xs.slice([splitIndex, 0], [-1, -1])
    ];
    const [trainYs, valYs] = [
        ys.slice([0, 0], [splitIndex, -1]),
        ys.slice([splitIndex, 0], [-1, -1])
    ];

    // Update UI with training parameters
    document.getElementById('total-epochs').textContent = epochs;
    document.getElementById('training-progress').style.display = 'block';

    // Training callback for progress updates
    const onEpochEnd = (epoch, logs) => {
        const progress = ((epoch + 1) / epochs) * 100;
        document.getElementById('training-progress-bar').style.width = `${progress}%`;
        document.getElementById('current-epoch').textContent = epoch + 1;
        document.getElementById('train-accuracy').textContent = logs.acc.toFixed(5);
        document.getElementById('valid-accuracy').textContent = logs.val_acc.toFixed(5);

        // Save training history
        tfTrainingHistory.accuracy.push(logs.acc);
        tfTrainingHistory.valAccuracy.push(logs.val_acc);
        tfTrainingHistory.loss.push(logs.loss);
        tfTrainingHistory.valLoss.push(logs.val_loss);

        // Update accuracy charts
        updateTrainingCharts(epoch + 1, logs);
    };

    // Start training
    console.log('Starting TensorFlow.js model training...');
    await tfModel.fit(trainXs, trainYs, {
        epochs,
        batchSize,
        validationData: [valXs, valYs],
        callbacks: { onEpochEnd }
    });

    // Clean up tensors to free memory
    xs.dispose();
    ys.dispose();
    trainXs.dispose();
    valXs.dispose();
    trainYs.dispose();
    valYs.dispose();

    // Finalize training UI
    finalizeTraining();
    console.log('TensorFlow.js model training completed');
}

/**
 * Update training accuracy charts during training
 * @param {number} epoch - Current epoch number
 * @param {Object} logs - Training metrics (accuracy, loss)
 */
function updateTrainingCharts(epoch, logs) {
    // Update training accuracy chart
    if (window.accuracyChart) {
        window.accuracyChart.data.labels.push(epoch);
        window.accuracyChart.data.datasets[0].data.push(logs.acc);
        window.accuracyChart.update();
    }

    // Update validation accuracy chart
    if (window.valAccuracyChart) {
        window.valAccuracyChart.data.labels.push(epoch);
        window.valAccuracyChart.data.datasets[0].data.push(logs.val_acc);
        window.valAccuracyChart.update();
    }
}

/**
 * Predict fraud probability for a single transaction
 * @param {Object} transaction - Transaction details (V1-V28, Time, Amount)
 * @returns {number} - Fraud probability (0-1)
 */
async function predictFraud(transaction) {
    if (!tfModel) {
        alert('Model not trained. Complete training first.');
        return 0;
    }

    // Prepare input tensor
    const input = tf.tensor2d([[
        transaction.time, transaction.v1, transaction.v2, transaction.v3, transaction.v4,
        transaction.v5, transaction.v6, transaction.v7, transaction.v8, transaction.v9,
        transaction.v10, transaction.v11, transaction.v12, transaction.v13, transaction.v14,
        transaction.v15, transaction.v16, transaction.v17, transaction.v18, transaction.v19,
        transaction.v20, transaction.v21, transaction.v22, transaction.v23, transaction.v24,
        transaction.v25, transaction.v26, transaction.v27, transaction.v28, transaction.amount
    ]]).div(tf.scalar(1000)); // Normalize input

    // Make prediction
    const prediction = await tfModel.predict(input).data();
    input.dispose(); // Clean up tensor

    return prediction[0]; // Return fraud probability (0-1)
}